package com.xhq.tools.rsa;

import com.sshtools.common.publickey.InvalidPassphraseException;
import com.sshtools.common.publickey.SshKeyUtils;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;

import java.io.IOException;
import java.io.StringReader;
import java.security.*;
import java.security.spec.EncodedKeySpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * @author huiny
 * @since 2021/8/14 22:35
 */
public class RsaKeyHelper {

    public static RsaDerKeyPair generateRSADerKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("RSA");
        keyPairGen.initialize(2048, new SecureRandom());
        KeyPair keyPair = keyPairGen.generateKeyPair();
        PrivateKey privateKey = keyPair.getPrivate();
        PublicKey publicKey = keyPair.getPublic();
        String publicKeyString = Base64Helper.encode(publicKey.getEncoded());
        String privateKeyString = Base64Helper.encode(privateKey.getEncoded());
        return new RsaDerKeyPair(publicKeyString, privateKeyString);
    }

    public static RsaDerKeyPair getDerKeyPairFromPemPrivateKey(String pemPrivateKey) throws IOException {
        try (StringReader keyReader = new StringReader(pemPrivateKey);
             PEMParser pemReader = new PEMParser(keyReader)) {
            JcaPEMKeyConverter converter = new JcaPEMKeyConverter();
            Object keyPair = pemReader.readObject();
            var pemKeyPair = ((PEMKeyPair) keyPair);
            PrivateKey privateKey = converter.getPrivateKey(pemKeyPair.getPrivateKeyInfo());
            PublicKey publicKey = converter.getPublicKey(pemKeyPair.getPublicKeyInfo());
            return new RsaDerKeyPair(Base64Helper.encode(publicKey.getEncoded()), Base64Helper.encode(privateKey.getEncoded()));
        }
    }

    public static RsaDerKeyPair getRsaDerKeyPairFromSshPrivateKey(String sshPrivateKey)
            throws IOException, InvalidPassphraseException {
        var SshKeyPair = SshKeyUtils.getPrivateKey(sshPrivateKey, null);
        PrivateKey privateKey = SshKeyPair.getPrivateKey().getJCEPrivateKey();
        PublicKey publicKey = SshKeyPair.getPublicKey().getJCEPublicKey();
        return new RsaDerKeyPair(Base64Helper.encode(publicKey.getEncoded()), Base64Helper.encode(privateKey.getEncoded()));
    }

    public static PrivateKey getPrivateKeyFromDerFile(byte[] keyBytes) throws Exception {
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory kf = KeyFactory.getInstance("RSA");
        return kf.generatePrivate(spec);
    }

    public static PrivateKey getPrivateKeyFromDerString(String derPrivateKey)
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(Base64Helper.decode(derPrivateKey));
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePrivate(keySpec);
    }

    public static PublicKey getPublicKeyFromDerFile(byte[] keyBytes) throws Exception {
        X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
        KeyFactory kf = KeyFactory.getInstance("RSA");
        return kf.generatePublic(spec);
    }

    public static PublicKey getPublicKeyFromBase64String(String publicKey)
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        EncodedKeySpec keySpec = new X509EncodedKeySpec(Base64Helper.decode(publicKey));
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePublic(keySpec);
    }
}
